import cv2
import numpy as np
import glob

def monocular_calibration(data_dir, board_size = (8,11), square_size = 10.0):

    images_dir = data_dir
    # 找棋盘格角点
    criteria_findCorner = (cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
    # 设置寻找亚像素角点的参数，采用的停止准则是最大循环次数30和最大误差容限0.001
    criteria_cornerSubPix = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001) # 阈值
    #棋盘格模板规格
    board_size = (8,11)
    square_size = 10.0  # 单位mm
    # 世界坐标系中的棋盘格点,例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
    objp = np.zeros((board_size[0]*board_size[1], 3), np.float32)
    objp[:,:2] = np.mgrid[0:board_size[1], 0:board_size[0]].T.reshape(-1,2) * square_size
    # self.objp = objp # 10 mm


    images = glob.glob(images_dir + "/*.bmp")
    # 储存棋盘格角点的世界坐标和图像坐标对
    objpoints = [] # 在世界坐标系中的三维点
    imgpoints = [] # 在图像平面的二维点
    i=0
    for fname in images:
        img = cv2.imread(fname)
        u, v = img.shape[:2]
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # 找到棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, (board_size[1], board_size[0]), criteria_findCorner)
        # 如果找到足够点对，将其存储起来
        if ret == True:
            print("i:", i)
            i = i+1
            # 在原角点的基础上寻找亚像素角点
            cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria_cornerSubPix)
            #追加进入世界三维点和平面二维点中
            objpoints.append(objp)
            imgpoints.append(corners)
            # 将角点在图像上显示
            cv2.drawChessboardCorners(img, (board_size[1], board_size[0]), corners, ret)
            cv2.namedWindow('findCorners', cv2.WINDOW_NORMAL)
            cv2.resizeWindow('findCorners', 640, 480)
            cv2.imshow('findCorners', img)
            cv2.waitKey(200)
        
    cv2.destroyAllWindows()
    # 标定
    print('正在计算')
    ret, cameraMatrix, distCoeffs, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
    print("ret:",ret)
    print("cameraMatrix:\n", cameraMatrix) # 内参数矩阵
    print("distCoeffs:\n", distCoeffs)   # 畸变系数   distortion cofficients = (k_1,k_2,p_1,p_2,k_3)
    
    print("正在写入文件")
    fs = cv2.FileStorage('camera_params.yaml', cv2.FILE_STORAGE_WRITE)
    fs.write("cameraMatrix", cameraMatrix)
    fs.write("distCoeffs", distCoeffs)
    fs.write("rvecs", np.array(rvecs))
    fs.write("tvecs", np.array(tvecs))
    fs.release()

    # newcameramtx, roi = cv2.getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, (u, v), 0, (u, v))
    # print('newcameramtx外参',newcameramtx)
    # print('roi:', roi)
    # #打开摄像机
    # camera=cv2.VideoCapture(0)
    # while True:
    #     (grabbed,frame)=camera.read()
    #     h1, w1 = frame.shape[:2]
    #     newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (u, v), 0, (u, v))
    #     # 纠正畸变
    #     dst1 = cv2.undistort(frame, mtx, dist, None, newcameramtx)
    #     #dst2 = cv2.undistort(frame, mtx, dist, None, newcameramtx)
    #     mapx,mapy=cv2.initUndistortRectifyMap(mtx,dist,None,newcameramtx,(w1,h1),5)
    #     dst2=cv2.remap(frame,mapx,mapy,cv2.INTER_LINEAR)
    #     # 裁剪图像，输出纠正畸变以后的图片
    #     x, y, w1, h1 = roi
    #     dst1 = dst1[y:y + h1, x:x + w1]

    #     #cv2.imshow('frame',dst2)
    #     #cv2.imshow('dst1',dst1)
    #     cv2.imshow('dst2', dst2)
    #     if cv2.waitKey(1) & 0xFF == ord('q'):  # 按q保存一张图片
    #         cv2.imwrite("../u4/frame.jpg", dst1)
    #         break

    # camera.release()
    # cv2.destroyAllWindows()

if __name__ == "__main__":
    # 可利用data/eyeinhand20250427路径下数据进行相机标定测试
    data_dir = "D:/VSCodeProjects/calibration/data/eyeinhand20250427"
    monocular_calibration(data_dir, board_size = (8,11), square_size = 10.0)